import os
import json
training_set_labels = []
testing_set_labels = []


DATASET_BASE_PATH = r"C:\Users\zhangbohang\Dataset"

def generate_labels(label_file_path):
    lines = None
    with open(label_file_path, "r") as f:
        lines = f.readlines()
    for line in lines:
        line_info_list = line.strip().split("\t")
        file_name = line_info_list[0]
        sheet_name = line_info_list[1]
        train_or_test = line_info_list[2]
        file_folder = line_info_list[3]
        
        dataset_name = file_folder.split("\\")[0]
        file_path = os.path.join(DATASET_BASE_PATH, file_folder, file_name)
        table_regions = line_info_list[4:]
        for table_region in table_regions:
            item = {
                "dataset_name": dataset_name,
                "file_path": file_path,
                "sheet_name": sheet_name,
                "table_region": table_region
            }
            if train_or_test == "training_set":
                training_set_labels.append(item)
            elif train_or_test == "testing_set":
                testing_set_labels.append(item)
            else:
                import pdb;pdb.set_trace()
    return training_set_labels, testing_set_labels

training_set_labels, testing_set_labels = generate_labels("dataset/Table range annotations.txt")
print(len(training_set_labels))
print(len(testing_set_labels))
training_labels_save_path = "labels/training_labels.txt"
testing_labels_save_path = "labels/testing_labels.txt"
with open(training_labels_save_path, "w") as f:
    f.write(json.dumps(training_set_labels))



